
open Common

module type POOLABLE_ITEM =
sig
  type t
  val create : unit -> t
end

module type POOLABLE =
sig

  type t 
  val create : unit -> t
  val init : int -> (int -> t)

end

module ArrayPoolable
  (T : sig type t val create : unit -> t end) : POOLABLE 
  with type t = T.t =
struct 

  module A = Array'

  type t = T.t

  let create = T.create

  let init n = 
    let default = T.create() in
    let arr = A.create n default in
    for i = 1 to n - 1 do
      arr.(i) <- T.create()
    done;
    fun n -> arr.(n)

end

module type T = 
sig
  
  module P : POOLABLE

  type t
  
  val proj : t -> P.t

  (** init size keep
      initializes the list, containing at most 'size' elements,
      and assuming you keep around 'keep' elements at any given time *)
  val init : size:int -> keep:int -> unit
  val used : unit -> int
  val size : unit -> int
  (** new_() returns a new memory location if one is available.  
      This is not the same as malloc, in the sense that if 
      do_gc() is called, this memory might be reclaimed.  Thus,
      do_gc must only be called when one is sure all necessary 
      values have been saved.  *)
  val new_ : unit -> t

  val save : t -> unit
  val delete : t -> unit

  val do_gc : unit -> unit

  module Ops : 
  sig
    val (!!) : t -> P.t
  end 

end 

module S(P : POOLABLE) : T with module P = P = 
struct
  
  module P = P
  module A = Array'
  module H = Hashtbl'

  type t = int

  let size = ref 0

  let (keepers : (int,unit) H.t ref) = ref (H.create 0)

  let (pool : (int -> P.t) ref) = ref (fun _ -> P.create())

  let (free_list : int array ref) = ref [||]

  let used = ref 0

  let proj n = (!pool) n

  let collect() = 
    let arr = !free_list in
    let ks = ref (sort (map fst (H.to_list !keepers))) in
    (** invariant: forall k \in ks, used <= k *)
    begin 
      used := 0;
      (** reset free_list *)
      for i = 0 to !size-1 do arr.(i) <- i done;
      (** put used cells at front, and swap unused cells backwards *)
      while !ks <> [] do
        let h,t = match !ks with h::t -> h,t | [] -> failwith "new_free_list" in 
        begin 
          A.swap arr !used h;
          incr used;
          ks := t;
        end
      done;
    end

  let do_gc() =
    (print_endline "collecting...";
     collect())

  let new_() = 
    if !used >= !size - 1 then failwith "new_: full" else
    let m = !free_list.(!used) in
      (incr used;m)

  let save m = 
    begin 
      H.unique_add !keepers m ();
      incr used;
    end

  let delete m = 
    if H.mem !keepers m then H.remove !keepers m 
    else failwith "free: not present"

  let called_init = ref false

  let init ~size:n ~keep = 
    let fl = A.create n 0 in
    begin 
      if !called_init then failwith "init: already initialized";
      print_endline ("initializing at size " ^ string_of_int n ^ "...");
      called_init := true;
      size := n;
      keepers := H.create keep;
      free_list := fl;
      for i = 0 to n-1 do fl.(i) <- i done;
      pool := P.init n;
    end

  let used() = !used

  let size() = !size

  module Ops =
  struct 
    let (!!) = proj
  end

end
